from datetime import datetime, timedelta
from decimal import Decimal, ROUND_DOWN

import pandas as pd
from fastapi import WebSocket

from app.api import env_parameter
from app.models.simulation import Simulation
from app.models.simulation_result import SimulationResult
from app.models.simulation_result_total import SimulationResultTotal
from app.models.simulation_result_total_map import SimulationResultTotalMap
from app.schemas.simulation_result import SimulationResultOut
from app.schemas.simulation_result_total import SimulationResultTotalCreate, SimulationResultTotalOutDetail, \
    SimulationResultTotalOut
from app.services.energy_dispatch.energy_management_input import EnergyManagementInput, EnergyInputSeries
from app.services.energy_dispatch.energy_dispatch_simple.energy_dispatch_simple import energy_management_strategy
from app.services.energy_dispatch.rl_energy_dispatch.dqn_agent import train


async def create_total(data: SimulationResultTotalCreate):
    obj = await SimulationResultTotal.create(**data.dict())
    return obj


async def get_total_by_id(total_id: int):
    return await SimulationResultTotal.get_or_none(simulation_result_total_id=total_id)


async def list_totals():
    return await SimulationResultTotal.all()


async def get_simulation_results_by_total_id(total_id: int):
    # 查总表记录
    total = await SimulationResultTotal.get_or_none(simulation_result_total_id=total_id)
    if not total:
        return None

    # 查总表关联的所有 simulation_result
    maps = await SimulationResultTotalMap.filter(simulation_result_total_id=total_id).prefetch_related("simulation_result")
    results = [m.simulation_result for m in maps]

    # 构造输出
    return SimulationResultTotalOutDetail(
        **SimulationResultTotalOut.from_orm(total).dict(),  # or use SimulationResultTotalOut.from_orm(total).dict()
        results=[SimulationResultOut.from_orm(r) for r in results]
    )

async def delete_total_by_id(total_id: int) -> int:
    """
    根据ID删除simulation_result_total及其关联的映射记录和simulation_result记录
    返回删除的simulation_result_total条数（0或1）
    """
    # 先查出所有关联的 simulation_result_id
    map_records = await SimulationResultTotalMap.filter(simulation_result_total_id=total_id).all()
    simulation_result_ids = [m.simulation_result_id for m in map_records]

    # 删除所有关联的 simulation_result 记录
    if simulation_result_ids:
        await SimulationResult.filter(simulation_result_id__in=simulation_result_ids).delete()

    # 删除映射表中相关记录
    await SimulationResultTotalMap.filter(simulation_result_total_id=total_id).delete()

    # 删除总表记录
    deleted_count = await SimulationResultTotal.filter(simulation_result_total_id=total_id).delete()

    return deleted_count

async def run_simulation(env_param_id: int, simulation_result_total_id: int) -> bool:
    """
    根据 env_param_id 运行一天（24小时）的储能调度仿真。
    保存每小时的 simulation_result，更新 simulation_result_total。
    """

    # 1. 获取输入数据
    data = await env_parameter.get_env_parameter_by_id(env_param_id)
    battery_info = data["battery"]
    prices = data["prices"]
    parameters = data["parameters"]
    granularity = data["granularity"]
    period = data["period"]
    initial_capacity = data["initial_capacity"]

    time_list = prices[0]["time_list"][:granularity * period]
    solar_price_values = next((p["value_list"] for p in prices if p["price_type"] == 0), [0.0] * 24)[:24]
    wind_price_values = next((p["value_list"] for p in prices if p["price_type"] == 1), [0.0] * 24)[:24]
    normal_price_values = next((p["value_list"] for p in prices if p["price_type"] == 2), [0.0] * 24)[:24]
    load_values = next((p["value_list"] for p in parameters if p["category_id"] == 5), [0.0] * 24)[:24]
    solar_values = next((p["value_list"] for p in parameters if p["category_id"] == 4), [0.0] * 24)[:24]
    wind_values = next((p["value_list"] for p in parameters if p["category_id"] == 3), [0.0] * 24)[:24]

    # 2. 初始化变量
    battery_energy = initial_capacity
    simulation_result_ids = []

    # 累加用字段
    total_cost = Decimal("0.0")
    total_emission = Decimal("0.0")
    solar_cost_total = Decimal("0.0")
    wind_cost_total = Decimal("0.0")
    grid_cost_total = Decimal("0.0")
    solar_emission_total = Decimal("0.0")
    wind_emission_total = Decimal("0.0")
    grid_emission_total = Decimal("0.0")
    solar_power_total = Decimal("0.0")
    wind_power_total = Decimal("0.0")
    grid_power_total = Decimal("0.0")
    total_power_total = Decimal("0.0")

    # 3. 每小时仿真
    for hour in range(24):
        current_time = time_list[hour]

        input_data = EnergyManagementInput(
            battery_energy=battery_energy,
            load=load_values[hour],
            solar_output=solar_values[hour],
            wind_output=wind_values[hour],
            price_grid=normal_price_values[hour],
            price_wind=wind_price_values[hour],
            price_solar=solar_price_values[hour],
            battery_capacity=battery_info["capacity"],
            battery_max_charge_power=battery_info["max_input_power"],
            battery_max_discharge_power=battery_info["max_output_power"],
        )

        result = energy_management_strategy(input_data, env_param_id, "Simple")

        # 更新电池状态
        battery_energy = result["battery_energy"]

        # 创建 simulation_result（每小时）
        sim = await SimulationResult.create(
            env_param_id=env_param_id,
            time_point=current_time,
            battery_usage_rate=result["battery_usage_rate"],
            battery_energy=battery_energy,
            grid_purchase=result["grid_purchase"],
            solar_supply=result["solar_supply"],
            wind_supply=result["wind_supply"],
            solar_power=result["solar_power"],
            grid_power=result["grid_power"],
            wind_power=result["wind_power"],
            total_power=result["total_power"],
            solar_cost=result["solar_cost"],
            wind_cost=result["wind_cost"],
            grid_cost=result["grid_cost"],
            solar_emission=result["solar_emission"],
            wind_emission=result["wind_emission"],
            grid_emission=result["grid_emission"],
        )
        simulation_result_ids.append(sim.simulation_result_id)

        # 累加汇总数据
        total_cost += result["total_cost"]
        total_emission += result["total_emission"]
        solar_cost_total += result["solar_cost"]
        wind_cost_total += result["wind_cost"]
        grid_cost_total += result["grid_cost"]
        solar_emission_total += result["solar_emission"]
        wind_emission_total += result["wind_emission"]
        grid_emission_total += result["grid_emission"]
        solar_power_total += result["solar_power"]
        wind_power_total += result["wind_power"]
        grid_power_total += result["grid_power"]
        total_power_total += result["total_power"]

    # 4. 更新汇总表 SimulationResultTotal
    await SimulationResultTotal.filter(simulation_result_total_id=simulation_result_total_id).update(
        env_param_id=env_param_id,
        battery_energy=battery_energy,
        solar_power=solar_power_total,
        grid_power=grid_power_total,
        wind_power=wind_power_total,
        total_power=total_power_total,
        solar_cost=solar_cost_total,
        wind_cost=wind_cost_total,
        grid_cost=grid_cost_total,
        solar_emission=solar_emission_total,
        wind_emission=wind_emission_total,
        grid_emission=grid_emission_total,
        period_unit="day",
        period=24
    )

    # 5. 建立映射关系
    for sim_id in simulation_result_ids:
        await SimulationResultTotalMap.create(
            simulation_result_total_id=simulation_result_total_id,
            simulation_result_id=sim_id
        )

    return True

def generate_24_hour_datetime_list() -> list[datetime]:
    base_time = datetime.today().replace(hour=0, minute=0, second=0, microsecond=0)
    return [base_time + timedelta(hours=i) for i in range(24)]

async def websocket_send_json(websocket:WebSocket, sim_dict: dict, simulation_type:str,simulation_id):
    # 把 datetime 转成字符串
    if isinstance(sim_dict.get("time_point"), datetime):
        sim_dict["time_point"] = sim_dict["time_point"].isoformat()

    for k, v in sim_dict.items():
        if isinstance(v, Decimal):
            sim_dict[k] = float(v)

    await websocket.send_json({
        "simulation_id": simulation_id,
        "simulation_type": simulation_type,
        "simulation_result": sim_dict,
    })

async def run_simulation_simple(ws:WebSocket,input_list:EnergyInputSeries,env_param_id: int,simulation_id:int,
                                battery_info,initial_capacity) -> bool:

    time_list = generate_24_hour_datetime_list()
    solar_price_values = input_list.solar_price_values
    wind_price_values = input_list.wind_price_values
    normal_price_values = input_list.normal_price_values
    load_values = input_list.load_values
    solar_values = input_list.solar_values
    wind_values = input_list.wind_values

    # 2. 初始化变量
    battery_energy = initial_capacity
    simulation_result_ids = []

    # 累加用字段
    total_cost = Decimal("0.0")
    total_emission = Decimal("0.0")
    solar_cost_total = Decimal("0.0")
    wind_cost_total = Decimal("0.0")
    grid_cost_total = Decimal("0.0")
    solar_emission_total = Decimal("0.0")
    wind_emission_total = Decimal("0.0")
    grid_emission_total = Decimal("0.0")
    solar_power_total = Decimal("0.0")
    wind_power_total = Decimal("0.0")
    grid_power_total = Decimal("0.0")
    total_power_total = Decimal("0.0")

    # 3. 每小时仿真
    for hour in range(24):
        current_time = time_list[hour]

        input_data = EnergyManagementInput(
            battery_energy=battery_energy,
            load=load_values[hour],
            solar_output=solar_values[hour],
            wind_output=wind_values[hour],
            price_grid=normal_price_values[hour],
            price_wind=wind_price_values[hour],
            price_solar=solar_price_values[hour],
            battery_capacity=battery_info["capacity"],
            battery_max_charge_power=battery_info["max_input_power"],
            battery_max_discharge_power=battery_info["max_output_power"],
        )


        result = energy_management_strategy(input_data, env_param_id, "simple")

        # 更新电池状态
        battery_energy = result["battery_energy"]

        # 创建 simulation_result（每小时）
        sim = await SimulationResult.create(
            env_param_id=env_param_id,
            time_point=current_time,
            battery_usage_rate=limit_decimal(result["battery_usage_rate"]),
            battery_energy=battery_energy,
            solar_power=limit_decimal(result["solar_power"]),
            grid_power=limit_decimal(result["grid_power"]),
            wind_power=limit_decimal(result["wind_power"]),
            total_power=limit_decimal(result["total_power"]),
            solar_cost=limit_decimal(result["solar_cost"]),
            wind_cost=limit_decimal(result["wind_cost"]),
            grid_cost=limit_decimal(result["grid_cost"]),
            solar_emission=limit_decimal(result["solar_emission"]),
            wind_emission=limit_decimal(result["wind_emission"]),
            grid_emission=limit_decimal(result["grid_emission"]),
        )
        simulation_result_ids.append(sim.simulation_result_id)


        sim_dict = sim.__dict__.copy()  # 复制一份防止修改原对象
        await websocket_send_json(ws, sim_dict,"simple",simulation_id)

        # 累加汇总数据
        total_cost += limit_decimal(result["total_cost"])
        total_emission += limit_decimal(result["total_emission"])
        solar_cost_total += limit_decimal(result["solar_cost"])
        wind_cost_total += limit_decimal(result["wind_cost"])
        grid_cost_total += limit_decimal(result["grid_cost"])
        solar_emission_total += limit_decimal(result["solar_emission"])
        wind_emission_total += limit_decimal(result["wind_emission"])
        grid_emission_total += limit_decimal(result["grid_emission"])
        solar_power_total += limit_decimal(result["solar_power"])
        wind_power_total += limit_decimal(result["wind_power"])
        grid_power_total += limit_decimal(result["grid_power"])
        total_power_total += limit_decimal(result["total_power"])

    # 4. 更新汇总表 SimulationResultTotal
    simTotal = await SimulationResultTotal.create(
        env_param_id=env_param_id,
        name="simple",
        strategy="simple",
        battery_energy=battery_energy,
        solar_power=solar_power_total,
        grid_power=grid_power_total,
        wind_power=wind_power_total,
        total_power=total_power_total,
        solar_cost=solar_cost_total,
        wind_cost=wind_cost_total,
        grid_cost=grid_cost_total,
        solar_emission=solar_emission_total,
        wind_emission=wind_emission_total,
        grid_emission=grid_emission_total,
        period_unit="day",
        period=24,
        simulation_id=simulation_id
    )

    # 5. 建立映射关系
    for sim_id in simulation_result_ids:
        await SimulationResultTotalMap.create(
            simulation_result_total_id=simTotal.simulation_result_total_id,
            simulation_result_id=sim_id
        )

    return True

def limit_decimal(val):
    return Decimal(str(val)).quantize(Decimal("0.0001"), rounding=ROUND_DOWN)

async def save_training_results(ws:WebSocket,env_param_id: int, episodes: int,
                                best_dqn_stats: dict,
                                best_dqn_actions: pd.DataFrame,
                                simulation_id:int,
                                strategy:str):
    """
    保存训练统计指标和动作记录到数据库
    """
    simulation_result_ids = []
    # 初始化累计变量
    total_battery_energy = Decimal('0')
    solar_power_total = Decimal('0')
    grid_power_total = Decimal('0')
    wind_power_total = Decimal('0')
    total_power_total = Decimal('0')

    solar_cost_total = Decimal('0')
    wind_cost_total = Decimal('0')
    grid_cost_total = Decimal('0')

    solar_emission_total = Decimal('0')
    wind_emission_total = Decimal('0')
    grid_emission_total = Decimal('0')

    time_list = generate_24_hour_datetime_list()

    # 1. 遍历每步动作，保存 SimulationResult 并累加数据
    for i, (_, row) in enumerate(best_dqn_actions.iterrows()):
        battery_energy = limit_decimal(row.get("battery_energy", 0))
        solar_power = limit_decimal(row.get("solar_power", 0))
        grid_power = limit_decimal(row.get("grid_power", 0))
        wind_power = limit_decimal(row.get("wind_power", 0))
        total_power = limit_decimal(row.get("total_power", 0))

        solar_cost = limit_decimal(row.get("solar_cost", 0))
        wind_cost = limit_decimal(row.get("wind_cost", 0))
        grid_cost = limit_decimal(row.get("grid_cost", 0))

        solar_emission = limit_decimal(row.get("solar_emission", 0))
        wind_emission = limit_decimal(row.get("wind_emission", 0))
        grid_emission = limit_decimal(row.get("grid_emission", 0))

        battery_usage_rate = limit_decimal(row.get("battery_usage_rate", 0))


        # 保存单步结果
        sim = await SimulationResult.create(
            env_param_id=env_param_id,
            time_point=time_list[i],
            battery_usage_rate=battery_usage_rate,
            battery_energy=battery_energy,
            solar_power=solar_power,
            grid_power=grid_power,
            wind_power=wind_power,
            total_power=total_power,
            solar_cost=solar_cost,
            wind_cost=wind_cost,
            grid_cost=grid_cost,
            solar_emission=solar_emission,
            wind_emission=wind_emission,
            grid_emission=grid_emission,
        )
        simulation_result_ids.append(sim.simulation_result_id)
        sim_dict = sim.__dict__.copy()  # 复制一份防止修改原对象
        await websocket_send_json(ws, sim_dict,strategy,simulation_id)

        # 累加汇总数据
        total_battery_energy += battery_energy
        solar_power_total += solar_power
        grid_power_total += grid_power
        wind_power_total += wind_power
        total_power_total += total_power

        solar_cost_total += solar_cost
        wind_cost_total += wind_cost
        grid_cost_total += grid_cost

        solar_emission_total += solar_emission
        wind_emission_total += wind_emission
        grid_emission_total += grid_emission

    # 2. 保存汇总表 SimulationResultTotal
    simTotal = await SimulationResultTotal.create(
        name=strategy,
        strategy=strategy,
        env_param_id=env_param_id,
        battery_energy = Decimal(str(best_dqn_actions.iloc[-1]["battery_energy"])) if len(best_dqn_actions) > 0 else Decimal('0'),
        solar_power=solar_power_total,
        grid_power=grid_power_total,
        wind_power=wind_power_total,
        total_power=total_power_total,
        solar_cost=solar_cost_total,
        wind_cost=wind_cost_total,
        grid_cost=grid_cost_total,
        solar_emission=solar_emission_total,
        wind_emission=wind_emission_total,
        grid_emission=grid_emission_total,
        period_unit="day",
        period=24,
        simulation_id=simulation_id
    )
    # 5. 建立映射关系
    for sim_id in simulation_result_ids:
        await SimulationResultTotalMap.create(
            simulation_result_total_id=simTotal.simulation_result_total_id,
            simulation_result_id=sim_id
        )
    print("训练结果和动作数据已保存到数据库。")